# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/8/25 上午11:27
# @Author  : shutao
# @FileName: quantum.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import paddle_quantum
import paddle
import numpy as np
from paddle_quantum.circuit import UAnsatz
from paddle_quantum.state import vec_random, vec

from paddle_quantum.utils import plot_state_in_bloch_sphere, plot_rotation_in_bloch_sphere


def single_qubit():
    """
    单量子比特
    """
    np.random.seed(42)
    num_samples = 15
    state = []

    for i in range(num_samples):
        # 创建单比特量子的线路
        cir = UAnsatz(1)
        # 生成随机的旋转角度
        phi, theta, omega = 2 * np.pi * np.random.uniform(size=3)
        phi = paddle.to_tensor(phi, dtype='float64')
        theta = paddle.to_tensor(theta, dtype='float64')
        omega = paddle.to_tensor(omega, dtype='float64')

        # 对单比特量子态做随机的旋转操作
        cir.rz(phi, 0)
        cir.ry(theta, 0)
        cir.rz(omega, 0)

        # 将线路的末态存储起来
        state.append(cir.run_state_vector())

    # 调用布洛赫球面展示功能， 输入state参数，并展示向量
    plot_state_in_bloch_sphere(state, show_arrow=True)


def locus_single_qubit():
    """
    单量子比特在布洛赫球面运行轨迹
    """
    # 创建单量子比特
    cir = UAnsatz(1)

    # 设置量子态的初始状态
    cir.ry(paddle.to_tensor(np.pi / 4, dtype="float64"), 0)
    cir.rz(paddle.to_tensor(np.pi / 2, dtype="float64"), 0)
    init_state = cir.run_density_matrix()

    # 将要进行的酉算子操作
    theta = np.pi / 2
    phi = 7 * np.pi / 8
    lam = 2 * np.pi
    rotating_angle = [theta, phi, lam]

    # 调用布洛赫球面展示功能，输入初始状态init_state，酉算子的旋转角度rotating_angle
    plot_rotation_in_bloch_sphere(init_state, rotating_angle)


def pauli_X():
    """
    X 单量子比特门
    """
    # 设置角度参数 theta = pi
    theta = np.array([np.pi])
    theta = paddle.to_tensor(theta)

    # 设置计算所需的量子比特数量
    num_qubits = 1

    # 初始化我们的单量子比特电路
    cir = UAnsatz(num_qubits)

    which_qubit = 0
    cir.rx(theta, which_qubit)

    # 打印出这个量子门
    print("量子门的矩阵表达式为：\n{}".format(cir.U.numpy()))


def pauli_Y():
    """
    Y 单比特量子门
    """
    theta = paddle.to_tensor(np.array([np.pi]))

    num_qubits = 1
    cir = UAnsatz(num_qubits)

    which_qubit = 0
    cir.ry(theta, which_qubit)

    # 打印出这个量子门
    print("量子门的矩阵表达式为：\n{}".format(cir.U.numpy()))


def TB_QNN():
    r"""
    多比特量子神经网络

    .. 三比特量子电路::
        --Ry(3.142)----*----Ry(3.142)---------------
               |
        --Ry(3.142)----x--------*--------Ry(3.142)--
                                |
        --Ry(3.142)-------------x--------Ry(3.142)--
    """
    # 设置角度参数
    theta = np.full([6], np.pi)
    theta = paddle.to_tensor(theta)

    # 初始化量子电路
    num_qubits = 3
    cir = UAnsatz(num_qubits)

    # 添加单比特旋转门
    cir.ry(theta[0], 0)
    cir.ry(theta[1], 1)
    cir.ry(theta[2], 2)

    # 添加两比特量子门
    cir.cnot([0, 1])
    cir.cnot([1, 2])

    # 添加单比特旋转门
    cir.ry(theta[3], 0)
    cir.ry(theta[4], 1)
    cir.ry(theta[5], 2)

    print('图中量子神经网络 U(theta=pi) 的电路表达式是:')
    print(cir)


def wave_function():
    """
    波函数向量表示
    复数向量表示和存储量子态
    """
    N = 20          # 设置量子比特数
    DEPTH = 6       # 设置量子电路深度
    theta = np.random.randn(DEPTH, N, 1)

    # 调用内置的 |00..0> 初始态
    initial_state1 = vec(0, N)
    # 调用内置的随机量子态
    initial_state2 = vec_random(N)

    theta = paddle.to_tensor(theta)
    initial_state = paddle.to_tensor(initial_state1)

    # 初始化量子电路
    cir = UAnsatz(N)

    # 添加深度为 Depth 的实数强纠缠结构QNN {Ry+CNOT's}
    cir.real_entangled_layer(theta, DEPTH)

    # 制备输出态
    # 如果用户不输入初始量子态，默认初始为|00..0>
    final_state = cir.run_state_vector(initial_state)
    print(final_state.numpy())


if __name__ == '__main__':
    # single_qubit()
    # locus_single_qubit()
    # pauli_X()
    # pauli_Y()
    # TB_QNN()
    wave_function()
